package study.ai.langchain4j.support.langchain4j;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import study.ai.langchain4j.service.ChatAssistant;
import study.ai.langchain4j.service.ChatStreamAssistant;

@Slf4j
@Configuration
public class AiServiceConfig {

    @Autowired
    private ChatLanguageModel chatLanguageModel;
    @Autowired
    StreamingChatLanguageModel streamingChatLanguageModel;

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    @Bean
    public ChatAssistant chatAssistant(EmbeddingStore<TextSegment> embeddingStore) {
        return AiServices.builder(ChatAssistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10)) //
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))       //
                .chatLanguageModel(chatLanguageModel)
                .build();
    }

    @Bean
    public ChatStreamAssistant chatSteamAssistant() {
        return AiServices.builder(ChatStreamAssistant.class)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .build();
    }
}
